# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import gzip
import os
import struct
import logging

from collections import namedtuple
from coin.base.datetime_util import to_datetime

try:
  from coin.feed.fastfeed.native_wrapper import run_from_feed_cache
except (ImportError, ModuleNotFoundError):
  logging.warning('Unable to import native module.')


class RecordWriter:
  def __init__(self, filepath):
    self.filepath = filepath

    # Ensure directory exists.
    dirpath = os.path.dirname(filepath)
    if not os.path.isdir(dirpath):
      os.makedirs(dirpath, exist_ok=True)

    if filepath.endswith('.gz'):
      self._fp = gzip.GzipFile(filepath, 'wb', compresslevel=9)
    else:
      self._fp = open(filepath, 'wb')
    self._bytes_written = 0

  def __del__(self):
    self.close()

  @property
  def bytes_written(self):
    return self._bytes_written

  def close(self):
    if self._fp:
      self._fp.close()
      self._fp = None

  def write(self, msg):
    assert isinstance(msg, bytes)
    header = struct.pack('i', len(msg))
    self._fp.write(header)
    self._fp.write(msg)


_FeedCacheKey = namedtuple('FeedCacheKey', ['machine', 'sub_req', 'worker_id', 'date'])


# TODO(jaewon): Remove `date` from FeedCacheKey
class FeedCacheKey(_FeedCacheKey):
  def __new__(cls, machine, date, sub_req, worker_id=None):
    # Override __new__, not __init__
    # https://stackoverflow.com/questions/42385916/inheriting-from-a-namedtuple-base-class-python
    self = super(FeedCacheKey, cls).__new__(cls,
                                            machine=machine,
                                            sub_req=sub_req,
                                            worker_id=(worker_id or '1'),
                                            date=to_datetime(date).date())
    return self

  @property
  def date_str(self):
    return self.date.strftime('%Y%m%d')

  def get_cache_filepath(self, product, extension=None):
    extension = extension or ''
    return os.path.join(self.machine,
                        self.date_str,
                        self.sub_req.to_str(),
                        '%s.W%s%s' % (product.symbol, self.worker_id, extension))


class FeedCacheWriter:
  def __init__(self, cache_key, compress=True, feed_cache_dir=None):
    self._cache_key = cache_key
    self._feed_cache_dir = os.path.expanduser(feed_cache_dir or '~/data/feed_cache')
    self._writer = {}
    self._compress = compress

  def __del__(self):
    self.close()

  def close(self):
    for writer in self._writer.values():
      writer.close()
    self._writer = {}

  def _get_writer(self, product):
    if product not in self._writer:
      ext = '.gz' if self._compress else ''
      cache_filepath = self._cache_key.get_cache_filepath(product, ext)
      self._writer[product] = RecordWriter(os.path.join(self._feed_cache_dir, cache_filepath))
    return self._writer[product]

  def write_proto_msg(self, product, msgs):
    writer = self._get_writer(product)
    for msg in msgs:
      writer.write(msg.SerializeToString())
